package assignment01;

import nlp.util.Counter;

public class GoodTuringCounter<E> extends Counter<E> {
	
	private static final long serialVersionUID = -3211277830509705043L;
	private static int THRESHOLD = 5;
	private double[] N = new double[THRESHOLD + 2];
	private Double unknownCount;

	public void smooth() {
		double totalCount = super.totalCount();
		Counter<E> helper = new Counter<E>();
		N[0] = totalCount;
		for (E key : keySet()) {
			int c = (int) getCount(key);
			if (c < THRESHOLD + 1) {
				N[c] += 1;
				helper.setCount(key, c);
			} else if (c == THRESHOLD + 1) { // deal with N(c+1) case
				N[c] += 1;
				setCount(key, c);
			} else {
				setCount(key, c);
			}
		}
		for (int i = 0; i <= THRESHOLD + 1; i++) {
			System.out.println("N(" + i + ") is " + N[i]);
		}
		
		for (E key: helper.keySet()) {
			int c = (int)helper.getCount(key);
			if (c <= THRESHOLD) {
				double cStar = Math.log((c + 1) * N[c + 1] / N[c]
						- (c * (THRESHOLD + 1) * N[THRESHOLD + 1]) / N[1])
						- Math.log(1.0 - (THRESHOLD + 1) * N[THRESHOLD + 1] / N[1]);
				cStar = Math.pow(Math.E, cStar);
//				System.out.println("cStar is " + cStar);
				setCount(key, cStar);
			}
		}
		// zeroCount = entryCount[1] / totalCount; REMEMBER, it's double, not int!
		// now I realize how important unit test is.
		unknownCount = N[1];
		System.out.println("zeroCount is " + unknownCount);
		
		///////
		validate();
	}
	
	/**
	public void normalize() {
		 double totalCount = totalCount();
		    for (E key : keySet()) {
		      setCount(key, getCount(key) / totalCount);
		    }
		    validate();
	}
	*/
	
	private void validate() {
		double sum = 0.0;
		double epsilon = .00001;
		for (E key : keySet()) {
			sum += getCount(key);
		}
		sum += getZeroCount();
		if (Math.abs(sum - totalCount()) < epsilon ) {
			System.out.println("validated " + sum);
		} else {
			System.out.println("failed " + sum);
		}
	}
	
	public double totalCount() {
		return N[0];
	}
	
	public double getZeroCount() {
		if (unknownCount == null || Double.isInfinite(Math.log(unknownCount))) {
			unknownCount = N[1];
		}
		return unknownCount;
	}
}
